// chsocket.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "chsocket.h"

#pragma comment(lib, "ws2_32.lib")

int chsocket::m_iSelectTimeout = 0;


// ============================================================================
// Construction of chsocket class
// ============================================================================
chsocket::chsocket()
{
	m_hdListenSocket = INVALID_SOCKET;
	m_hdConnSocket = INVALID_SOCKET;
	set_listen_ip("127.0.0.1");
	m_iListenPort = 0;
	m_iConnPort = 0;
	m_ulConnIP = 0;
	m_iMaxSendBuff = 0;
	m_iMaxRecvBuff = 0;
	m_iSendBuffSize = 0;
	m_iRecvBuffSize = 0;
	m_pSendBuff = NULL;
	m_pRecvBuff = NULL;

	init_buff(CHSOCKET_SEND_BUFF_LEN, CHSOCKET_RECV_BUFF_LEN);
	set_msg(__FUNCTION__ " New instance: %X", this);
}


// ============================================================================
// Destruction of chsocket class
// ============================================================================
chsocket::~chsocket()
{
	set_msg(__FUNCTION__ " Release instance: %X", this);
	release_buff();
}


// ============================================================================
// Init winsock by calling WSAStartup()
// ============================================================================
int chsocket::init_winsock()
// return: <0> succeed, <-1> failed
{
	WSADATA wsaData;
	int iResult;
	iResult = WSAStartup(MAKEWORD(2,2), &wsaData);
    if (NO_ERROR != iResult)
	{
      set_msg(__FUNCTION__" Error WSAStartup: %d\n", iResult);
      return -1;
    }

	set_msg(__FUNCTION__" Winsock initialized.");

	return 0;
}


// ============================================================================
// Release winsock by calling WSACleanup()
// ============================================================================
void chsocket::release_winsock()
{
	WSACleanup();
	set_msg(__FUNCTION__" Winsock released.");
}


// ============================================================================
// Output tracing log or error message
// ============================================================================
void chsocket::set_msg(const char *pstrMsg, ...)
// [pstrMsg]: message to be output, support variable parameters as printf(...)
{
	if(NULL == pstrMsg)
	{
		return;
	}

	va_list argp;
	va_start(argp, pstrMsg);
	_vsnprintf_s(m_strMsg, CHSOCKET_MSG_LEN, pstrMsg, argp);
	printf("-- %s\n", m_strMsg);
}


// ============================================================================
// Get last message
// ============================================================================
const char* chsocket::get_msg()
// return: pointer to message address
{
	return m_strMsg;
}


// ============================================================================
// Allocate memory for send&recv buffer
// ============================================================================
int chsocket::init_buff(int iSendBuffSize, int iRecvBuffSize)
// [iSendBuffSize]: maximum size of send buffer, recommended value = 65535
// [iRecvBuffSize]: maximum size of recv buffer, recommended value = 65535
// return: <0> succeed, <-1> failed
{
	if(iSendBuffSize > 0)
	{
		// release first
		if(m_iMaxSendBuff > 0 && NULL != m_pSendBuff)
		{
			delete m_pSendBuff;
		}
		m_pSendBuff = NULL;
		m_iMaxSendBuff = 0;
		m_iSendBuffSize = 0;

		m_pSendBuff = new char[iSendBuffSize + 1];
		if(NULL == m_pSendBuff)
		{
			set_msg(__FUNCTION__ " Memory not enough for send buff. (LEN=%d)", iSendBuffSize);
			return -1;
		}

		m_iMaxSendBuff = iSendBuffSize;
	}

	if(iRecvBuffSize > 0)
	{
		// release first
		if(m_iMaxRecvBuff > 0 && NULL != m_pRecvBuff)
		{
			delete m_pRecvBuff;
		}
		m_pRecvBuff = NULL;
		m_iMaxRecvBuff = 0;
		m_iRecvBuffSize = 0;

		m_pRecvBuff = new char[iRecvBuffSize + 1];
		if(NULL == m_pRecvBuff)
		{
			set_msg(__FUNCTION__ " Memory not enough for recv buff. (LEN=%d)", iRecvBuffSize);
			return -1;
		}

		m_iMaxRecvBuff = iRecvBuffSize;
	}

	set_msg(__FUNCTION__ " Send buffer size:%d/%d, recv buffer size:%d/%d", 
		m_iSendBuffSize, m_iMaxSendBuff, m_iRecvBuffSize, m_iMaxRecvBuff);
	return 0;
}


// ============================================================================
// Release send&recv buffer
// ============================================================================
void chsocket::release_buff()
{
	if(m_iMaxSendBuff > 0 && NULL != m_pSendBuff)
	{
		delete m_pSendBuff;
	}
	m_pSendBuff = NULL;
	m_iMaxSendBuff = 0;
	m_iSendBuffSize = 0;

	if(m_iMaxRecvBuff > 0 && NULL != m_pRecvBuff)
	{
		delete m_pRecvBuff;
	}
	m_pRecvBuff = NULL;
	m_iMaxRecvBuff = 0;
	m_iRecvBuffSize = 0;

	set_msg(__FUNCTION__ " Send buffer size:%d/%d, recv buffer size:%d/%d", 
		m_iSendBuffSize, m_iMaxSendBuff, m_iRecvBuffSize, m_iMaxRecvBuff);
}


// ============================================================================
// Set synchronize mode as Block or Asynchronize
// ============================================================================
void chsocket::set_sync_mode(u_long iMode)
// [iMode]: <CHSOCKET_BLOCK_MODE:0> Block Mode, <CHSOCKET_ASYNC_MODE:1> Asynchronize Mode
{
	if(INVALID_SOCKET != m_hdListenSocket)
	{
		ioctlsocket(m_hdListenSocket, FIONBIO, &iMode);
		set_msg(__FUNCTION__ " Listen socket set sync mode %d. (SOCKET=%d)", iMode, m_hdListenSocket);
	}

	if(INVALID_SOCKET != m_hdConnSocket)
	{
		ioctlsocket(m_hdConnSocket, FIONBIO, &iMode);
		set_msg(__FUNCTION__ " Connect socket set sync mode %d. (SOCKET=%d)", iMode, m_hdConnSocket);
	}
}


// ============================================================================
// Set listen IP address
// ============================================================================
void chsocket::set_listen_ip(const char *pstrIP)
// [pstrIP]: IP string, e.g. "127.0.0.1"
{
	m_ulListenIP = inet_addr(pstrIP);
}


// ============================================================================
// Set listen port
// ============================================================================
void chsocket::set_listen_port(int iListenPort)
// [iListenPort]: port to be listened
{
	m_iListenPort = iListenPort;
}


// ============================================================================
// Listen with IP & port specified by set_listen_ip() & set_listen_port()
// ============================================================================
int	chsocket::listen()
// return: <0> succeed, <-1> failed
{
	m_hdListenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (m_hdListenSocket == INVALID_SOCKET)
	{
		set_msg(__FUNCTION__ " Error calling socket(). (errorcode=%ld)", WSAGetLastError());
	    return -1;
	}

	sockaddr_in service;
	service.sin_family = AF_INET;
	service.sin_addr.s_addr = m_ulListenIP;
	service.sin_port = htons(m_iListenPort);

	if (SOCKET_ERROR == ::bind(m_hdListenSocket, (SOCKADDR*) &service, sizeof(service)))
	{
		set_msg(__FUNCTION__ " Error calling bind(). (SOCKET=%d, IP=%d.%d.%d.%d PORT=%d)", m_hdListenSocket,
			m_ulListenIP&0xFF, (m_ulListenIP>>8)&0xFF, (m_ulListenIP>>16)&0xFF, (m_ulListenIP>>24)&0xFF, m_iListenPort);
		closesocket(m_hdListenSocket);
		m_hdListenSocket = INVALID_SOCKET;
		return -1;
	}

	if (SOCKET_ERROR == ::listen( m_hdListenSocket, SOMAXCONN ))
	{
		set_msg(__FUNCTION__  " Error calling listen(). (SOCKET=%d, IP=%d.%d.%d.%d PORT=%d)", m_hdListenSocket,
			m_ulListenIP&0xFF, (m_ulListenIP>>8)&0xFF, (m_ulListenIP>>16)&0xFF, (m_ulListenIP>>24)&0xFF, m_iListenPort);
		return -1;
	}

	set_msg(__FUNCTION__" Socket listened. (SOCKET=%d, IP=%d.%d.%d.%d PORT=%d)", m_hdListenSocket,
		m_ulListenIP&0xFF, (m_ulListenIP>>8)&0xFF, (m_ulListenIP>>16)&0xFF, (m_ulListenIP>>24)&0xFF, m_iListenPort);
	return 0;
}


// ============================================================================
// Close listen socket
// ============================================================================
int chsocket::close_listen()
// return: <0> succeed, <-1> failed
{
	if(INVALID_SOCKET == m_hdListenSocket)
	{
		set_msg(__FUNCTION__ " Listen socket not initlized.");
		return -1;
	}

	::closesocket(m_hdListenSocket);
	m_hdListenSocket = INVALID_SOCKET;
	set_msg(__FUNCTION__ " Listen socket %d closed.", m_hdListenSocket);
	return 0;
}


// ============================================================================
// Accept new connection to listen socket
// ============================================================================
int chsocket::accept(chsocket *pchSocket)
// [pchSocket]: <NULL> check connect queue only, <&chsocket> new connection instance
// return: <0> new connection arrived, <1> no new connection, <-1> failed 
{
	SOCKET hAcceptSocket;
	fd_set rset;
	timeval tv;

	FD_ZERO(&rset);
	FD_SET(m_hdListenSocket, &rset);
	tv.tv_sec = m_iSelectTimeout / 1000;
	tv.tv_usec = m_iSelectTimeout % 1000;

	int iRet = select(m_hdListenSocket + 1, &rset, NULL, NULL, &tv);

	// select() failed
	if(SOCKET_ERROR == iRet)
	{
		set_msg(__FUNCTION__ " Select failed: %d", WSAGetLastError());
		return -1;
	}

	// select() timeout
	if(0 == iRet)
	{
		return 1;
	}

	// if no client socket specified, return 0 without calling accept()
	if(NULL == pchSocket)
	{
		return 0;
	}

	// new connection in the queue, accept the connection.
	hAcceptSocket = ::accept(m_hdListenSocket, NULL, NULL);
	if (INVALID_SOCKET == hAcceptSocket)
	{
		set_msg(__FUNCTION__ " Accept failed: %d", WSAGetLastError());
		return -1;
	}

	if(0 != pchSocket->attach_conn_sock(hAcceptSocket))
	{
		return -1;
	}

    set_msg(__FUNCTION__ " Client connected. (SOCKET=%d)", hAcceptSocket);
	return 0;
}



// ============================================================================
// Set connect IP address
// ============================================================================
void chsocket::set_conn_ip(const char *pstrIP)
// [pstrIP]: IP string, e.g. "127.0.0.1"
{
	m_ulConnIP = inet_addr(pstrIP);
}


// ============================================================================
// Set connect port
// ============================================================================
void chsocket::set_conn_port(int iConnPort)
// [iListenPort]: port to be listened
{
	m_iConnPort = iConnPort;
}


// ============================================================================
// Connect to IP & port specified by set_conn_ip() & set_conn_port()
// ============================================================================
int chsocket::connect()
// return: <0> succeed, <-1> failed
{
	if(NULL != m_hdConnSocket)
	{
		close_connect();
	}

	m_hdConnSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (m_hdConnSocket == INVALID_SOCKET)
	{
		set_msg(__FUNCTION__ " Error calling socket(). (errorcode=%ld)", WSAGetLastError());
	    return -1;
	}

	sockaddr_in clientService; 
	clientService.sin_family = AF_INET;
	clientService.sin_addr.s_addr = m_ulConnIP;
	clientService.sin_port = htons(m_iConnPort);

	if (SOCKET_ERROR == ::connect(m_hdConnSocket, (SOCKADDR*) &clientService, sizeof(clientService)))
	{
		set_msg(__FUNCTION__ " Error calling connect(). (SOCKET=%d, IP=%d.%d.%d.%d PORT=%d)", m_hdConnSocket,
			m_ulConnIP&0xFF, (m_ulConnIP>>8)&0xFF, (m_ulConnIP>>16)&0xFF, (m_ulConnIP>>24)&0xFF, m_iConnPort);
		closesocket(m_hdConnSocket);
		m_hdConnSocket = INVALID_SOCKET;
		return -1;
	}

	set_msg(__FUNCTION__" Connected to server. (SOCKET=%d, IP=%d.%d.%d.%d PORT=%d)", m_hdConnSocket,
		m_ulConnIP&0xFF, (m_ulConnIP>>8)&0xFF, (m_ulConnIP>>16)&0xFF, (m_ulConnIP>>24)&0xFF, m_iConnPort);

	return 0;
}


// ============================================================================
// Close connect socket
// ============================================================================
int chsocket::close_connect()
// return: <0> succeed, <-1> failed
{
	if(INVALID_SOCKET == m_hdConnSocket)
	{
		set_msg(__FUNCTION__ " Connect socket not initlized.");
		return -1;
	}

	::closesocket(m_hdConnSocket);
	m_hdConnSocket = INVALID_SOCKET;
	set_msg(__FUNCTION__ " Connect socket %d closed.", m_hdConnSocket);
	return 0;
}


// ============================================================================
// Attach existing socket handle as connect socket
// ============================================================================
int	chsocket::attach_conn_sock(SOCKET hConnSocket)
// [hConnSocket]: SOCKET handle
// return: <0> succeed, <-1> failed
{
	if(INVALID_SOCKET != m_hdConnSocket)
	{
		close_connect();
	}

	if(INVALID_SOCKET == hConnSocket)
	{
		set_msg(__FUNCTION__ " Invalid socket specified.", hConnSocket);
		return -1;
	}

	m_hdConnSocket = hConnSocket;
	return 0;
}


// ============================================================================
// Send data to connect socket
// ============================================================================
int	chsocket::send(const char* pstrSend, int iSize)
// [pstrSend]: data to be sent
// [iSize]: size of data
// return: <0> succeed, <-1> failed
{
	if(INVALID_SOCKET == m_hdConnSocket)
	{
		set_msg(__FUNCTION__ " Socket not connected.");
		return -1;
	}

	// send buffer first 
	if(iSize + m_iSendBuffSize > m_iMaxSendBuff)
	{
		send_buffer();
	}

	// buffer is full
	if(iSize + m_iSendBuffSize > m_iMaxSendBuff)
	{
		set_msg(__FUNCTION__ " Send buffer overflow. (%d+%d>%d)", 
			iSize, m_iSendBuffSize, m_iMaxSendBuff);
		return -1;
	}

	memcpy(m_pSendBuff + m_iSendBuffSize, pstrSend, iSize);
	m_iSendBuffSize += iSize;
	return send_buffer();
}


// ============================================================================
// Send buffer to socket
// ============================================================================
int	chsocket::send_buffer()
// return: <0> succeed, <-1> failed
{
	if(NULL == m_pSendBuff || m_iMaxSendBuff < m_iSendBuffSize)
	{
		set_msg(__FUNCTION__ " Error send buffer! (buff=%X, %d/%d)", 
			m_pSendBuff, m_iSendBuffSize, m_iMaxSendBuff);
		return -1;
	}

	if(m_iSendBuffSize == 0)
	{
		return 0;
	}

	int iSent = ::send(m_hdConnSocket, m_pSendBuff, m_iSendBuffSize, 0);
	if(SOCKET_ERROR == iSent)
	{
		set_msg(__FUNCTION__ " Send error %d. (SOCKET=%d)", WSAGetLastError(), m_hdConnSocket);
		return -1;
	}

	if(iSent >= 0)
	{
		if(iSent < m_iSendBuffSize)
		{
			memcpy(m_pSendBuff, m_pSendBuff + iSent, m_iSendBuffSize - iSent);
			m_iSendBuffSize -= iSent;
		}
		else
		{
			m_iSendBuffSize = 0;
		}

		set_msg(__FUNCTION__ " %d bytes sent. (SOCKET=%d)", iSent, m_hdConnSocket);
	}
	return 0;
}


// ============================================================================
// Recv data from connect socket
// ============================================================================
int	chsocket::recv()
// return: <>=0> bytes received, <-1> failed
{
	if(NULL == m_pRecvBuff || m_iMaxRecvBuff < m_iRecvBuffSize)
	{
		set_msg(__FUNCTION__ " Error recv buffer! (buff=%X, %d/%d)", 
			m_pRecvBuff, m_iRecvBuffSize, m_iMaxRecvBuff);
		return -1;
	}

	//  check recv buffer
	fd_set fds_red;
	timeval tv;

	FD_ZERO(&fds_red);
	FD_SET(m_hdListenSocket, &fds_red);
	tv.tv_sec = m_iSelectTimeout / 1000;
	tv.tv_usec = m_iSelectTimeout % 1000;
	int iRet = select(m_hdListenSocket + 1, &fds_red, NULL, NULL, &tv);

	// recv buffer is empty
	if(0 == FD_ISSET(m_hdListenSocket, &fds_red))
	{
		return 0;
	}

	int iRecv = ::recv(m_hdConnSocket, m_pRecvBuff, m_iMaxRecvBuff - m_iRecvBuffSize, 0);
	if(iRecv < 0)
	{
		set_msg(__FUNCTION__ " Error recv! (SOCKET=%d)", m_hdConnSocket);
		return -1;
	}
	if(iRecv == 0)
	{
		set_msg(__FUNCTION__ " Socket closed! (SOCKET=%d)", m_hdConnSocket);
		return -1;
	}

	m_iRecvBuffSize += iRecv;
	m_pRecvBuff[m_iRecvBuffSize] = 0;
	set_msg(__FUNCTION__ " %d bytes received. (SOCKET=%d, BUFF=%d)", 
		iRecv, m_hdConnSocket, m_iRecvBuffSize);

	return iRecv;
}


// ============================================================================
// Get recv buffer address
// ============================================================================
const char* chsocket::get_recv_buff()
// return: pointer to recv buffer
{
	return m_pRecvBuff;
}